# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/28 11:06 
@Author  : ZhangShenao 
@File    : 1.使用RunnableParallel实现并行链.py
@Desc    : 使用RunnableParallel实现并行链
RunnableParallel是LangChain封装的支持运行多个Runnable的类
它一般用于操作Runnable的输出,以匹配序列中下一个Runnable的输入
起到并行运行Runnable,并格式化输出结构的作用
"""
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel

# 编排两个prompt
joke_prompt = ChatPromptTemplate.from_template('请给我讲一个关于{subject}的冷笑话')
poem_prompt = ChatPromptTemplate.from_template('请帮我写一首关于{subject}的诗，尽可能简单一点')

# 创建智普AI LLM
llm = ChatZhipuAI(model="GLM-4-Air")

# 创建OutputParser
parser = StrOutputParser()

# 构建两个Chain
joke_chain = joke_prompt | llm | parser
poem_chain = poem_prompt | llm | parser

# 使用RunnableParallel,将两个Chain并行运行
parallel_chain = RunnableParallel(joke=joke_chain, poem=poem_chain)

# 运行RunnableParallel,多个Chain会并行执行,并返回一个最终结果的dict
result_dict = parallel_chain.invoke({'subject': '狗'})

# 打印结果
print(result_dict)
